import torch
from torch.autograd import Function
from ..box_utils import decode, nms, center_size
# from data import prior_box_config as cfg
import time
class Timer(object):
    """A simple timer."""
    def __init__(self):
        self.total_time = 0.
        self.calls = 0
        self.start_time = 0.
        self.diff = 0.
        self.average_time = 0.

    def tic(self):
        # using time.time instead of time.clock because time time.clock
        # does not normalize for multithreading
        self.start_time = time.time()

    def toc(self, average=True):
        self.diff = time.time() - self.start_time
        self.total_time += self.diff
        self.calls += 1
        self.average_time = self.total_time / self.calls
        if average:
            return self.average_time
        else:
            return self.diff

class Detection_RefineDet(object):
    """At test time, Detect is the final layer of SSD.  Decode location preds,
    apply non-maximum suppression to location predictions based on conf
    scores and threshold to a top_k number of output predictions for both
    confidence score and locations.
    """
    def __init__(self, num_classes, top_k, conf_thresh, nms_thresh, objectness_thresh,
                keep_top_k, cfg):
        self.num_classes = num_classes
        self.top_k = top_k
        self.keep_top_k = keep_top_k
        self.nms_thresh = nms_thresh
        if nms_thresh <= 0:
            raise ValueError('nms_threshold must be non negative.')  
        self.conf_thresh = conf_thresh
        self.objectness_thresh = objectness_thresh
        self.variance = cfg['variance']
    # @staticmethod
    def forward(self, arm_loc_data, arm_conf_data, odm_loc_data,odm_conf_data,
                prior_data):
        """
        Args:
            loc_data: (tensor) Loc preds from loc layers
                Shape: [batch,num_priors,4]
            conf_data: (tensor) Shape: Conf preds from conf layers
                Shape: [batch,num_priors,num_classes]
            prior_data: (tensor) Prior boxes and variances from priorbox layers
                Shape: [num_priors,4]
        """
        _t = Timer()
        _t.tic()
        loc_data = odm_loc_data
        conf_data = odm_conf_data

        arm_object_conf = arm_conf_data.data[:,:,1:]    # 提取arm输出有目标那一类的概率
        no_object_index = arm_object_conf <= self.objectness_thresh   # conf小于预制的位置会被设置为true
        conf_data[no_object_index.expand_as(conf_data)] = 0 # 将odm输出内容中背景预选框所有预测概率设为0

        batch_size = loc_data.size(0)
        num_priors = prior_data.size(0)
        output = torch.zeros(batch_size, self.num_classes, self.top_k, 5)
        conf_pred = conf_data.transpose(2,1) # [batch_size, num_classes, num_priors]
        
        # Decode predictions into bboxes.
        for i in range(batch_size):
            default = decode(arm_loc_data[i], prior_data, self.variance)    # decode结果为两个对角坐标
            default = center_size(default)
            decoded_boxes = decode(loc_data[i], default, self.variance)
            conf_scores = conf_pred[i].clone()  # [num_classes, num_priors]
            # 对每一类目标进行nms（非极大值抑制）
            for cl in range(1, self.num_classes):   # 从1开始，不考虑0（背景）
                c_mask = conf_scores[cl].gt(self.conf_thresh)   # 生成和conf_scores[cl]相同的张量，其中大于conf_thresh的位置为true
                scores = conf_scores[cl][c_mask]    # [num_mask]
                if scores.size(0) == 0:
                    continue    # 没有检测到此类别的框，跳过
                l_mash = c_mask.unsqueeze(1).expand_as(decoded_boxes)   # [num_priors, 4]
                boxes = decoded_boxes[l_mash].view(-1,4)    # [num_mask, 4]
                # 返回每一类中置信度最高且没有重叠的框，置信度阈值为nms_thresh，取超过预制的前top_k个
                ids, count = nms(boxes, scores, self.nms_thresh, self.top_k)
                output[i, cl, :count] = torch.cat((scores[ids[:count]].unsqueeze(1),
                                                    boxes[ids[:count]]), 1) # [count, 5(置信度，坐标)]
        '''
        下段代码似乎不起作用
        '''
        flt = output.contiguous().view(batch_size, -1, 5)
        _, idx = flt[:, :, 0].sort(1, descending=True)  # 将候选框的置信度由高向低排
        _, rank = idx.sort(1)   # 获得原顺序随机的候选框按照从高到低顺序排列后的位置
        flt[(rank < self.keep_top_k).unsqueeze(-1).expand_as(flt)].fill_(0)
        detection_refinedet_time = _t.toc()
        print('detection_refinedet_time:{:.3f}'.format(detection_refinedet_time))
        return output


